package com.kingwang.study.spring.security.demo.security.repository;

import com.kingwang.study.spring.security.demo.security.constant.SecurityConstants;
import com.kingwang.study.spring.security.demo.security.token.TokenBasedAuthenticationResultToken;
import com.kingwang.study.spring.security.demo.security.token.session.TokenSession;
import com.kingwang.study.spring.security.demo.security.token.session.TokenSessionManager;
import lombok.Data;
import org.apache.commons.lang3.StringUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.web.context.HttpRequestResponseHolder;
import org.springframework.security.web.context.SecurityContextRepository;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Token认证凭据的缓存存储和还原
 * <p>
 * 用于从认证凭据缓存或Token会话管理器中还原回Spring Security的SecurityContext和其中的Authentication对象
 * 使Spring Security认为认证成功
 */
public class HttpTokenSecurityContextRepository implements SecurityContextRepository {
    private Map<String, SecurityContext> repository = new ConcurrentHashMap<>();

    private TokenSessionManager tokenSessionManager;

    public HttpTokenSecurityContextRepository(TokenSessionManager tokenSessionManager) {
        this.tokenSessionManager = tokenSessionManager;
    }

    /**
     * 加载SecurityContext
     * 注意：虽然 {@link SecurityContextRepository#loadContext(HttpRequestResponseHolder)} 的注释信息中说明不要返回null
     * 但是由于我们是同时支持Session和Token登录，后续在 {@link SecurityContextCompositeRepository} 中还有其它的 {@link SecurityContextRepository}
     * 进行处理，所以返回null
     * <p>
     * 使用单个 {@link SecurityContextRepository} 的情况下绝对不能返回null！
     *
     * @param requestResponseHolder
     * @return
     */
    @Override
    public SecurityContext loadContext(HttpRequestResponseHolder requestResponseHolder) {
        HttpServletRequest request = requestResponseHolder.getRequest();

        // 从HttpServletRequest中获取Token
        AuthorizationToken token = getToken(request);

        if (token == null || !token.isValid()) {
            /**
             * 如果Token不存在或格式非法，且该URL为Token的登录URL，生成一个空 SecurityContext
             * 目的是不让后面的 {@link org.springframework.security.web.context.HttpSessionSecurityContextRepository} 接管 SecurityContext的生成
             * 因为它里面有特殊的逻辑对HttpServletResponse做了封装（副作用）
             * 接管后，HttpSessionSecurityContextRepository#saveContext(SecurityContext, HttpServletRequest,
             * HttpServletResponse) 方法就不会将SecurityContext保存到 HttpSessionSecurityContextRepository 中
             */
            if (isTokenLoginUrl(request)) {
                return SecurityContextHolder.createEmptyContext();
            } else {
                return null;
            }

//            return SecurityContextHolder.createEmptyContext();
        }

        /**
         * 先从内存缓存中获取SecurityContext
         * 如果缓存中不存在，从Token会话管理器中还原回SecurityContext
         */
        SecurityContext result = this.repository.get(token.getSessionId());

        if (result == null && request.getSession(false) == null) {
            result = SecurityContextHolder.createEmptyContext();

            TokenSession session = this.tokenSessionManager.getSession(token.getSessionId());
            if (session != null) {
                result.setAuthentication(new TokenBasedAuthenticationResultToken(session.getUsername(), null,
                        session.getType(), session.getToken(), session.getUserDetails().getAuthorities()));
            }
        }

        return result;
    }

    /**
     * 保存SecurityContext
     * 由于在Token会话管理器中已经存储了Token的信息，所以这里只需要存储到内存缓存中，以加快认证效率
     *
     * @param securityContext
     * @param request
     * @param response
     */
    @Override
    public void saveContext(SecurityContext securityContext, HttpServletRequest request, HttpServletResponse response) {
        AuthorizationToken token = getToken(request);

        if (token != null) {
            if (securityContext == null) {
                this.repository.remove(token.getSessionId());
            } else {
                this.repository.put(token.getSessionId(), securityContext);
            }
        }
    }

    /**
     * 检查本Repository中是否包含SecurityContext
     *
     * @param request
     * @return
     */
    @Override
    public boolean containsContext(HttpServletRequest request) {
        AuthorizationToken token = getToken(request);

        if (token == null) {
            return isTokenLoginUrl(request);
        }

        boolean result = this.repository.containsKey(token.getSessionId());
        if (!result) {
            TokenSession session = this.tokenSessionManager.getSession(token.getSessionId());

            return session != null;
        }

        return result;
    }

    private boolean isTokenLoginUrl(HttpServletRequest request) {
        return request.getRequestURI().equalsIgnoreCase(SecurityConstants.URL_LOGIN_TOKEN);
    }

    private AuthorizationToken getToken(HttpServletRequest request) {
        String authorizationHeader = request.getHeader(HttpHeaders.AUTHORIZATION);

        if (StringUtils.isBlank(authorizationHeader)) {
            return null;
        }

        String[] tokenParts = authorizationHeader.split(" ");

        AuthorizationToken result = new AuthorizationToken();

        if (tokenParts.length > 1) {
            result.setTokenType(tokenParts[0]);
            result.setToken(tokenParts[1]);
        }

        return result;
    }

    @Data
    private class AuthorizationToken {
        private String tokenType;

        private String token;

        public String getSessionId() {
            return HttpTokenSecurityContextRepository.this.tokenSessionManager.generateSessionId(this.tokenType, this.token);
        }

        public boolean isValid() {
            return StringUtils.isNotBlank(this.tokenType) && StringUtils.isNotBlank(this.token);
        }
    }
}
